{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at hfl/chinese-roberta-wwm-ext were not used when initializing BertForSequenceClassification: ['bert.encoder.layer.2.intermediate.dense.bias', 'bert.encoder.layer.6.output.LayerNorm.weight', 'bert.encoder.layer.6.attention.self.query.weight', 'bert.encoder.layer.7.attention.self.key.bias', 'bert.encoder.layer.5.output.LayerNorm.bias', 'bert.encoder.layer.4.intermediate.dense.weight', 'bert.encoder.layer.11.attention.self.value.bias', 'cls.predictions.transform.LayerNorm.bias', 'bert.encoder.layer.9.attention.output.dense.weight', 'bert.encoder.layer.5.attention.output.dense.weight', 'bert.encoder.layer.9.output.dense.bias', 'bert.encoder.layer.6.output.LayerNorm.bias', 'bert.encoder.layer.10.intermediate.dense.bias', 'bert.encoder.layer.4.attention.output.dense.weight', 'bert.encoder.layer.9.attention.output.LayerNorm.weight', 'bert.encoder.layer.10.attention.self.query.weight', 'bert.encoder.layer.8.intermediate.dense.bias', 'cls.predictions.transform.LayerNorm.weight', 'bert.encoder.layer.10.attention.output.LayerNorm.bias', 'bert.encoder.layer.7.output.dense.weight', 'bert.encoder.layer.10.output.LayerNorm.weight', 'bert.encoder.layer.7.attention.output.dense.bias', 'bert.encoder.layer.3.intermediate.dense.bias', 'bert.encoder.layer.7.intermediate.dense.bias', 'bert.encoder.layer.11.output.dense.weight', 'cls.seq_relationship.weight', 'bert.encoder.layer.6.output.dense.weight', 'bert.encoder.layer.3.attention.output.LayerNorm.bias', 'bert.encoder.layer.10.intermediate.dense.weight', 'bert.encoder.layer.8.output.dense.weight', 'bert.encoder.layer.8.attention.self.query.weight', 'bert.encoder.layer.9.attention.self.query.bias', 'bert.encoder.layer.2.attention.self.value.bias', 'bert.encoder.layer.2.output.LayerNorm.weight', 'cls.seq_relationship.bias', 'bert.encoder.layer.9.attention.output.dense.bias', 'bert.encoder.layer.3.attention.self.key.weight', 'bert.encoder.layer.3.attention.self.value.weight', 'bert.encoder.layer.7.output.LayerNorm.bias', 'cls.predictions.decoder.weight', 'cls.predictions.bias', 'bert.encoder.layer.4.attention.output.dense.bias', 'bert.encoder.layer.3.attention.self.value.bias', 'bert.encoder.layer.7.attention.self.key.weight', 'bert.encoder.layer.9.attention.self.value.weight', 'bert.encoder.layer.3.output.dense.bias', 'bert.encoder.layer.3.attention.self.query.weight', 'bert.encoder.layer.3.attention.output.dense.bias', 'bert.encoder.layer.4.intermediate.dense.bias', 'bert.encoder.layer.8.attention.self.key.weight', 'bert.encoder.layer.8.attention.self.query.bias', 'bert.encoder.layer.10.attention.self.query.bias', 'bert.encoder.layer.5.attention.self.value.bias', 'bert.encoder.layer.10.attention.self.value.bias', 'bert.encoder.layer.8.attention.self.value.weight', 'bert.encoder.layer.7.attention.output.dense.weight', 'bert.encoder.layer.6.attention.output.dense.bias', 'bert.encoder.layer.5.attention.self.query.weight', 'bert.encoder.layer.7.attention.self.query.bias', 'bert.encoder.layer.9.intermediate.dense.weight', 'bert.encoder.layer.5.attention.self.key.weight', 'bert.encoder.layer.6.attention.self.key.weight', 'bert.encoder.layer.3.output.LayerNorm.bias', 'bert.encoder.layer.4.attention.self.key.bias', 'bert.encoder.layer.4.output.dense.weight', 'bert.encoder.layer.6.attention.output.dense.weight', 'bert.encoder.layer.6.attention.self.value.weight', 'bert.encoder.layer.5.intermediate.dense.weight', 'bert.encoder.layer.3.output.dense.weight', 'bert.encoder.layer.9.intermediate.dense.bias', 'bert.encoder.layer.10.attention.output.LayerNorm.weight', 'bert.encoder.layer.3.attention.self.key.bias', 'bert.encoder.layer.10.attention.self.value.weight', 'bert.encoder.layer.2.intermediate.dense.weight', 'bert.encoder.layer.3.attention.output.LayerNorm.weight', 'bert.encoder.layer.5.output.dense.weight', 'bert.encoder.layer.8.attention.output.dense.bias', 'bert.encoder.layer.6.attention.self.query.bias', 'bert.encoder.layer.5.attention.self.query.bias', 'bert.encoder.layer.9.attention.self.query.weight', 'bert.encoder.layer.6.intermediate.dense.weight', 'bert.encoder.layer.11.attention.output.dense.bias', 'bert.encoder.layer.11.intermediate.dense.bias', 'bert.encoder.layer.2.attention.output.dense.bias', 'bert.encoder.layer.2.output.dense.bias', 'bert.encoder.layer.4.output.LayerNorm.weight', 'bert.encoder.layer.8.attention.output.dense.weight', 'bert.encoder.layer.11.attention.self.key.weight', 'bert.encoder.layer.7.attention.self.value.bias', 'bert.encoder.layer.11.attention.self.value.weight', 'bert.encoder.layer.4.attention.self.query.weight', 'bert.encoder.layer.8.intermediate.dense.weight', 'bert.encoder.layer.2.output.LayerNorm.bias', 'bert.encoder.layer.6.attention.self.key.bias', 'bert.encoder.layer.2.attention.output.LayerNorm.bias', 'bert.encoder.layer.10.attention.output.dense.bias', 'bert.encoder.layer.4.attention.output.LayerNorm.bias', 'bert.encoder.layer.11.output.LayerNorm.weight', 'bert.encoder.layer.11.attention.self.query.bias', 'bert.encoder.layer.7.intermediate.dense.weight', 'bert.encoder.layer.5.attention.self.value.weight', 'bert.encoder.layer.6.intermediate.dense.bias', 'bert.encoder.layer.6.output.dense.bias', 'bert.encoder.layer.9.output.dense.weight', 'bert.encoder.layer.10.output.dense.bias', 'bert.encoder.layer.2.attention.output.dense.weight', 'bert.encoder.layer.6.attention.self.value.bias', 'bert.encoder.layer.2.attention.self.key.bias', 'bert.encoder.layer.9.attention.self.key.weight', 'bert.encoder.layer.11.attention.self.query.weight', 'bert.encoder.layer.7.attention.self.query.weight', 'bert.encoder.layer.11.attention.output.LayerNorm.weight', 'bert.encoder.layer.3.attention.self.query.bias', 'bert.encoder.layer.10.attention.output.dense.weight', 'bert.encoder.layer.8.output.LayerNorm.bias', 'bert.encoder.layer.9.output.LayerNorm.weight', 'bert.encoder.layer.11.output.dense.bias', 'bert.encoder.layer.5.intermediate.dense.bias', 'bert.encoder.layer.2.output.dense.weight', 'bert.encoder.layer.8.attention.output.LayerNorm.bias', 'bert.encoder.layer.10.attention.self.key.weight', 'bert.encoder.layer.4.attention.self.query.bias', 'bert.encoder.layer.8.attention.self.value.bias', 'bert.encoder.layer.5.attention.self.key.bias', 'bert.encoder.layer.7.output.dense.bias', 'bert.encoder.layer.9.attention.self.key.bias', 'cls.predictions.transform.dense.weight', 'bert.encoder.layer.10.output.dense.weight', 'bert.encoder.layer.2.attention.self.value.weight', 'bert.encoder.layer.2.attention.self.query.bias', 'bert.encoder.layer.3.output.LayerNorm.weight', 'bert.encoder.layer.4.attention.self.value.bias', 'bert.encoder.layer.7.output.LayerNorm.weight', 'bert.encoder.layer.10.output.LayerNorm.bias', 'bert.encoder.layer.11.attention.output.LayerNorm.bias', 'bert.encoder.layer.3.attention.output.dense.weight', 'bert.encoder.layer.5.attention.output.LayerNorm.bias', 'bert.encoder.layer.8.output.LayerNorm.weight', 'cls.predictions.transform.dense.bias', 'bert.encoder.layer.2.attention.output.LayerNorm.weight', 'bert.encoder.layer.7.attention.output.LayerNorm.weight', 'bert.encoder.layer.7.attention.output.LayerNorm.bias', 'bert.encoder.layer.4.attention.self.value.weight', 'bert.encoder.layer.6.attention.output.LayerNorm.weight', 'bert.encoder.layer.5.attention.output.dense.bias', 'bert.encoder.layer.11.output.LayerNorm.bias', 'bert.encoder.layer.6.attention.output.LayerNorm.bias', 'bert.encoder.layer.5.output.dense.bias', 'bert.encoder.layer.7.attention.self.value.weight', 'bert.encoder.layer.4.output.dense.bias', 'bert.encoder.layer.8.output.dense.bias', 'bert.encoder.layer.2.attention.self.key.weight', 'bert.encoder.layer.9.output.LayerNorm.bias', 'bert.encoder.layer.8.attention.self.key.bias', 'bert.encoder.layer.9.attention.output.LayerNorm.bias', 'bert.encoder.layer.11.attention.output.dense.weight', 'bert.encoder.layer.11.attention.self.key.bias', 'bert.encoder.layer.4.attention.self.key.weight', 'bert.encoder.layer.5.attention.output.LayerNorm.weight', 'bert.encoder.layer.9.attention.self.value.bias', 'bert.encoder.layer.3.intermediate.dense.weight', 'bert.encoder.layer.4.output.LayerNorm.bias', 'bert.encoder.layer.5.output.LayerNorm.weight', 'bert.encoder.layer.8.attention.output.LayerNorm.weight', 'bert.encoder.layer.11.intermediate.dense.weight', 'bert.encoder.layer.2.attention.self.query.weight', 'bert.encoder.layer.10.attention.self.key.bias', 'bert.encoder.layer.4.attention.output.LayerNorm.weight']\n",
      "- This IS expected if you are initializing BertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at hfl/chinese-roberta-wwm-ext and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from transformers import BertTokenizer, BertForSequenceClassification, AdamW, BertConfig\n",
    "from torch.optim.lr_scheduler import StepLR\n",
    "\n",
    "# 自定义数据集类\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, data_path, tokenizer):\n",
    "        self.sentences, self.labels = self.load_data(data_path)\n",
    "        self.tokenizer = tokenizer\n",
    "\n",
    "    def load_data(self, data_path):\n",
    "        sentences = []\n",
    "        labels = []\n",
    "        with open(data_path, 'r', encoding='utf-8') as file:\n",
    "            for line in file:\n",
    "                sentence, label = line.strip().split('&')\n",
    "                sentences.append(sentence)\n",
    "                labels.append(int(label))\n",
    "        return sentences, labels\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.sentences)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        sentence = self.sentences[idx]\n",
    "        label = self.labels[idx]\n",
    "\n",
    "        # 使用BERT的tokenizer对句子进行编码\n",
    "        encoding = self.tokenizer.encode_plus(\n",
    "            sentence,\n",
    "            add_special_tokens=True,\n",
    "            padding='max_length',\n",
    "            truncation=True,\n",
    "            max_length=128,\n",
    "            return_tensors='pt'\n",
    "        )\n",
    "\n",
    "        input_ids = encoding['input_ids'].squeeze()\n",
    "        attention_mask = encoding['attention_mask'].squeeze()\n",
    "\n",
    "        return {\n",
    "            'input_ids': input_ids,\n",
    "            'attention_mask': attention_mask,\n",
    "            'label': torch.tensor(label)\n",
    "        }\n",
    "\n",
    "\n",
    "# 设置训练参数\n",
    "data_path = './data.txt'\n",
    "model_name = 'hfl/chinese-roberta-wwm-ext'\n",
    "batch_size = 32\n",
    "label_cnt = 7\n",
    "learning_rate = 5e-6\n",
    "num_epochs = 27\n",
    "\n",
    "# 加载预训练的BERT模型和tokenizer\n",
    "tokenizer = BertTokenizer.from_pretrained(model_name)\n",
    "config = BertConfig.from_pretrained(model_name)\n",
    "config.num_hidden_layers = 2\n",
    "config.num_labels = label_cnt   # 将 num_labels 添加到 config 中\n",
    "model = BertForSequenceClassification.from_pretrained(model_name, config=config)\n",
    "# model.classifier.add_module('dropout', torch.nn.Dropout(p=0.5))\n",
    "# model.classifier.add_module('fc1', torch.nn.Linear(768, 256))  # 在原有的分类器上增加一个全连接层\n",
    "# model.classifier.add_module('fc2', torch.nn.Linear(256, 128)) # 再加一层\n",
    "# model.classifier.add_module('fc3', torch.nn.Linear(128, 7))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\12236\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\transformers\\optimization.py:407: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "# 加载数据集\n",
    "dataset = CustomDataset(data_path, tokenizer)\n",
    "\n",
    "# 划分训练集和验证集\n",
    "train_size = int(0.8 * len(dataset))\n",
    "valid_size = len(dataset) - train_size\n",
    "train_dataset, valid_dataset = torch.utils.data.random_split(\n",
    "    dataset, [train_size, valid_size])\n",
    "\n",
    "# 创建数据加载器\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "# 定义优化器和损失函数\n",
    "# weight_decay为L2正则化 防止过拟合\n",
    "# scheduler 配置学习率衰减\n",
    "optimizer = AdamW(model.parameters(), lr=learning_rate, weight_decay=0.0001)\n",
    "scheduler = StepLR(optimizer, step_size=18, gamma=0.1)\n",
    "loss_fn = torch.nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "if torch.cuda.device_count() > 1:\n",
    "    model = torch.DataParallel(model, device_ids=[0, 1, 2])\n",
    "\n",
    "model.to(device)\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/27 - Loss: 1.7449 - Valid Loss: 1.5058 - Accuracy: 0.3494 - Valid Accuracy: 0.4929\n",
      "Epoch 2/27 - Loss: 1.1859 - Valid Loss: 0.8354 - Accuracy: 0.5962 - Valid Accuracy: 0.7333\n",
      "Epoch 3/27 - Loss: 0.6783 - Valid Loss: 0.5053 - Accuracy: 0.8048 - Valid Accuracy: 0.8333\n",
      "Epoch 4/27 - Loss: 0.4465 - Valid Loss: 0.3816 - Accuracy: 0.8765 - Valid Accuracy: 0.8778\n",
      "Epoch 5/27 - Loss: 0.3399 - Valid Loss: 0.3306 - Accuracy: 0.9065 - Valid Accuracy: 0.8873\n",
      "Epoch 6/27 - Loss: 0.2820 - Valid Loss: 0.2957 - Accuracy: 0.9226 - Valid Accuracy: 0.8984\n",
      "Epoch 7/27 - Loss: 0.2483 - Valid Loss: 0.2709 - Accuracy: 0.9267 - Valid Accuracy: 0.9040\n",
      "Epoch 8/27 - Loss: 0.2220 - Valid Loss: 0.2568 - Accuracy: 0.9347 - Valid Accuracy: 0.9079\n",
      "Epoch 9/27 - Loss: 0.1912 - Valid Loss: 0.2465 - Accuracy: 0.9436 - Valid Accuracy: 0.9127\n",
      "Epoch 10/27 - Loss: 0.1809 - Valid Loss: 0.2448 - Accuracy: 0.9454 - Valid Accuracy: 0.9127\n",
      "Epoch 11/27 - Loss: 0.1630 - Valid Loss: 0.2322 - Accuracy: 0.9510 - Valid Accuracy: 0.9175\n",
      "Epoch 12/27 - Loss: 0.1438 - Valid Loss: 0.2258 - Accuracy: 0.9575 - Valid Accuracy: 0.9206\n",
      "Epoch 13/27 - Loss: 0.1324 - Valid Loss: 0.2243 - Accuracy: 0.9633 - Valid Accuracy: 0.9230\n",
      "Epoch 14/27 - Loss: 0.1202 - Valid Loss: 0.2220 - Accuracy: 0.9635 - Valid Accuracy: 0.9143\n",
      "Epoch 15/27 - Loss: 0.1085 - Valid Loss: 0.2242 - Accuracy: 0.9684 - Valid Accuracy: 0.9278\n",
      "Epoch 16/27 - Loss: 0.0973 - Valid Loss: 0.2178 - Accuracy: 0.9730 - Valid Accuracy: 0.9238\n",
      "Epoch 17/27 - Loss: 0.0929 - Valid Loss: 0.2218 - Accuracy: 0.9732 - Valid Accuracy: 0.9286\n",
      "Epoch 18/27 - Loss: 0.0860 - Valid Loss: 0.2285 - Accuracy: 0.9770 - Valid Accuracy: 0.9222\n",
      "Epoch 19/27 - Loss: 0.0743 - Valid Loss: 0.2185 - Accuracy: 0.9803 - Valid Accuracy: 0.9270\n",
      "Epoch 20/27 - Loss: 0.0743 - Valid Loss: 0.2181 - Accuracy: 0.9790 - Valid Accuracy: 0.9286\n",
      "Epoch 21/27 - Loss: 0.0695 - Valid Loss: 0.2201 - Accuracy: 0.9837 - Valid Accuracy: 0.9270\n",
      "Epoch 22/27 - Loss: 0.0714 - Valid Loss: 0.2190 - Accuracy: 0.9807 - Valid Accuracy: 0.9286\n",
      "Epoch 23/27 - Loss: 0.0728 - Valid Loss: 0.2197 - Accuracy: 0.9819 - Valid Accuracy: 0.9310\n",
      "Epoch 24/27 - Loss: 0.0732 - Valid Loss: 0.2202 - Accuracy: 0.9801 - Valid Accuracy: 0.9286\n",
      "Epoch 25/27 - Loss: 0.0706 - Valid Loss: 0.2232 - Accuracy: 0.9821 - Valid Accuracy: 0.9294\n",
      "Epoch 26/27 - Loss: 0.0726 - Valid Loss: 0.2246 - Accuracy: 0.9807 - Valid Accuracy: 0.9278\n",
      "Epoch 27/27 - Loss: 0.0693 - Valid Loss: 0.2280 - Accuracy: 0.9811 - Valid Accuracy: 0.9310\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    correct_predictions = 0\n",
    "    total_predictions = 0\n",
    "\n",
    "    for batch in train_loader:\n",
    "        input_ids = batch['input_ids'].to(device)\n",
    "        attention_mask = batch['attention_mask'].to(device)\n",
    "        labels = batch['label'].to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)\n",
    "        loss = outputs.loss\n",
    "        total_loss += loss.item()\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        _, predicted_labels = torch.max(outputs.logits, dim=1)\n",
    "        correct_predictions += (predicted_labels == labels).sum().item()\n",
    "        total_predictions += labels.size(0)\n",
    "    \n",
    "    # 在验证集上评估模型\n",
    "    model.eval()\n",
    "    valid_loss = 0\n",
    "    correct_valid_predictions = 0\n",
    "    total_valid_predictions = 0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for batch in valid_loader:\n",
    "            input_ids = batch['input_ids'].to(device)\n",
    "            attention_mask = batch['attention_mask'].to(device)\n",
    "            labels = batch['label'].to(device)\n",
    "\n",
    "            outputs = model(input_ids, attention_mask=attention_mask)\n",
    "            logits = outputs.logits\n",
    "            _, predicted_labels = torch.max(logits, dim=1)\n",
    "\n",
    "            valid_loss += loss_fn(logits, labels).item()\n",
    "            correct_valid_predictions += (predicted_labels == labels).sum().item()\n",
    "            total_valid_predictions += labels.size(0)\n",
    "\n",
    "    epoch_loss = total_loss / len(train_loader)\n",
    "    epoch_valid_loss = valid_loss / len(valid_loader)\n",
    "    accuracy = correct_predictions / total_predictions\n",
    "    valid_accuracy = correct_valid_predictions / total_valid_predictions\n",
    "    scheduler.step()\n",
    "\n",
    "    print(f'Epoch {epoch + 1}/{num_epochs} - Loss: {epoch_loss:.4f} - Valid Loss: {epoch_valid_loss:.4f} - Accuracy: {accuracy:.4f} - Valid Accuracy: {valid_accuracy:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.9780848022868033\n"
     ]
    }
   ],
   "source": [
    "# 计算整体准确率\n",
    "v_loader = DataLoader(dataset)\n",
    "model.eval()\n",
    "correct_predictions = 0\n",
    "total_predictions = 0\n",
    "with torch.no_grad():\n",
    "        for batch in v_loader:\n",
    "            input_ids = batch['input_ids'].to(device)\n",
    "            attention_mask = batch['attention_mask'].to(device)\n",
    "            labels = batch['label'].to(device)\n",
    "\n",
    "            outputs = model(input_ids, attention_mask=attention_mask)\n",
    "            logits = outputs.logits\n",
    "            _, predicted_labels = torch.max(logits, dim=1)\n",
    "\n",
    "            correct_predictions += (predicted_labels == labels).sum().item()\n",
    "            total_predictions += labels.size(0)\n",
    "\n",
    "valid_accuracy = correct_predictions / total_predictions\n",
    "\n",
    "print(valid_accuracy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 保存模型\n",
    "torch.save(model.state_dict(), 'Roberta3.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "输入句子: 2016.9-2019.7          广州二中           高中\n",
      "预测标签: 0\n"
     ]
    }
   ],
   "source": [
    "def predict_label(sentence, model, tokenizer):\n",
    "    # 对输入句子进行编码\n",
    "    encoding = tokenizer.encode_plus(\n",
    "        sentence,\n",
    "        add_special_tokens=True,\n",
    "        padding='max_length',\n",
    "        truncation=True,\n",
    "        max_length=128,\n",
    "        return_tensors='pt'\n",
    "    )\n",
    "\n",
    "    input_ids = encoding['input_ids'].squeeze()\n",
    "    attention_mask = encoding['attention_mask'].squeeze()\n",
    "\n",
    "    # 在模型中进行前向传播\n",
    "    with torch.no_grad():\n",
    "        outputs = model(input_ids.unsqueeze(0), attention_mask=attention_mask.unsqueeze(0))\n",
    "        logits = outputs.logits\n",
    "\n",
    "    # 获取预测的标签\n",
    "    _, predicted_label = torch.max(logits, dim=1)\n",
    "\n",
    "    return predicted_label.item()\n",
    "\n",
    "# # 加载训练好的模型参数\n",
    "model.load_state_dict(torch.load('Roberta3.pt'))\n",
    "\n",
    "# 设置模型为评估模式\n",
    "model.eval()\n",
    "\n",
    "# 测试例子\n",
    "input_sentence = '2016.9-2019.7          广州二中           高中'\n",
    "predicted_label = predict_label(input_sentence, model, tokenizer)\n",
    "\n",
    "print(f'输入句子: {input_sentence}')\n",
    "print(f'预测标签: {predicted_label}')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "黄岩新前卫生院 -- 工作/项目经历\n",
      "广州市第三人民医院 -- 工作/项目经历\n",
      "广州医科大学 -- 基本信息\n",
      "黄岩新前卫生院 -- 工作/项目经历\n",
      "广州铁塔第一医院 -- 工作/项目经历\n",
      "北京协和医院 -- 工作/项目经历\n",
      "上海交通大学医学院附属瑞金医院 -- 工作/项目经历\n",
      "广州中山大学附属第一医院 -- 工作/项目经历\n",
      "成都华西医院 -- 工作/项目经历\n",
      "武汉大学人民医院 -- 工作/项目经历\n",
      "复旦大学附属华山医院 -- 工作/项目经历\n",
      "中山大学附属第三医院 -- 工作/项目经历\n",
      "南京大学医学院附属鼓楼医院 -- 工作/项目经历\n",
      "重庆医科大学附属第一医院 -- 工作/项目经历\n",
      "浙江大学医学院附属第一医院 -- 工作/项目经历\n",
      "2009.09-2013.06 复旦大学 环境科学 -- 基本信息\n",
      "2013.09-2017.06 吉林大学 旅游管理 -- 基本信息\n",
      "4-6k -- 垃圾信息\n",
      "8-10k -- 垃圾信息\n",
      "1-3k -- 垃圾信息\n",
      "4-5w -- 垃圾信息\n",
      "1-7u -- 垃圾信息\n",
      "6-9p -- 垃圾信息\n",
      "许爱礼 -- 基本信息\n",
      "陈虹荣 -- 基本信息\n",
      "陈虹荣 求职意向：工商管理 -- 基本信息\n",
      "胡健 求职意向：健身教练 -- 求职意向\n",
      "33 -- 垃圾信息\n",
      "34 -- 垃圾信息\n",
      "林竹水 -- 基本信息\n",
      "王爱乐 -- 基本信息\n",
      "2017.06-至今       安宁家乐福超市 -- 工作/项目经历\n",
      "2015.09-2015.12\t 兰州大学食堂咖啡店 -- 工作/项目经历\n",
      "实习护士 -- 基本信息\n",
      "临床护士 -- 基本信息\n",
      "管理学院 -- 垃圾信息\n"
     ]
    }
   ],
   "source": [
    "label = {\n",
    "    0: '基本信息',\n",
    "    1: '求职意向',\n",
    "    2: '工作/项目经历',\n",
    "    3: '获得奖项',\n",
    "    4: '个人能力',\n",
    "    5: '垃圾信息',\n",
    "    6: '重要标注'\n",
    "}\n",
    "sentences = [\n",
    "    '黄岩新前卫生院',\n",
    "    '广州市第三人民医院',\n",
    "    '广州医科大学',\n",
    "    '黄岩新前卫生院',\n",
    "    '广州铁塔第一医院',\n",
    "    '北京协和医院',\n",
    "    '上海交通大学医学院附属瑞金医院',\n",
    "    '广州中山大学附属第一医院',\n",
    "    '成都华西医院',\n",
    "    '武汉大学人民医院',\n",
    "    '复旦大学附属华山医院',\n",
    "    '中山大学附属第三医院',\n",
    "    '南京大学医学院附属鼓楼医院',\n",
    "    '重庆医科大学附属第一医院',\n",
    "    '浙江大学医学院附属第一医院',\n",
    "    '2009.09-2013.06 复旦大学 环境科学',\n",
    "    '2013.09-2017.06 吉林大学 旅游管理',\n",
    "    '4-6k',\n",
    "    '8-10k',\n",
    "    '1-3k',\n",
    "    '4-5w',\n",
    "    '1-7u',\n",
    "    '6-9p',\n",
    "    '许爱礼',\n",
    "    '陈虹荣',\n",
    "    '陈虹荣 求职意向：工商管理',\n",
    "    '胡健 求职意向：健身教练',\n",
    "    '33',\n",
    "    '34',\n",
    "    '林竹水',\n",
    "    '王爱乐',\n",
    "    '2017.06-至今       安宁家乐福超市',\n",
    "    '2015.09-2015.12\t 兰州大学食堂咖啡店',\n",
    "    '实习护士',\n",
    "    '临床护士',\n",
    "    '管理学院',\n",
    "    '韩健毓'\n",
    "]\n",
    "\n",
    "for s in sentences:\n",
    "    p = predict_label(s, model, tokenizer)\n",
    "    print(f'{s} -- {label[p]}')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
